package com.aispeech.asr.clib.utils;

import com.aispeech.asr.clib.decoder.AudioDecoder;
import com.aispeech.asr.clib.decoder.EAudioCodec;
import com.aispeech.asr.clib.vad.VadSplitter;
import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;

import java.io.IOException;
import java.io.InputStream;

public class AudioSplitUtils {
    private final static int BUFF_LEN = 3200;

    public static SplittedPcmData decode(int decodeSampleRate, InputStream stream, EAudioCodec codec, int sampleRate, int channel
            , int sampleBytes) throws IOException {
        return split(decodeSampleRate, stream, codec, sampleRate, channel, sampleBytes, true, null, 0);
    }

    public static SplittedPcmData split(int decodeSampleRate, InputStream stream, EAudioCodec codec, int sampleRate, int channel
            , int sampleBytes, String vadResPath, int vadPauseTime, int maxSpeechTime) throws IOException {
        return split(decodeSampleRate, stream, codec, sampleRate, channel, sampleBytes, false, vadResPath, vadPauseTime, maxSpeechTime);
    }

    public static SplittedPcmData split(int decodeSampleRate, InputStream stream, EAudioCodec codec, int sampleRate,
                                        int channel, int sampleBytes, boolean skipVad, String vadResPath,
                                        int vadPauseTime) throws IOException {
        return split(decodeSampleRate, stream, codec, sampleRate, channel, sampleBytes, skipVad, vadResPath, vadPauseTime, 30);
    }

    public static SplittedPcmData split(int decodeSampleRate, InputStream stream, EAudioCodec codec, int sampleRate,
                                        int channel, int sampleBytes, boolean skipVad, String vadResPath, int vadPauseTime, int maxSpeechTime) throws IOException {
        SplittedPcmData splittedPcmData = new SplittedPcmData();
        AudioDecoder audioDecoder = new AudioDecoder(decodeSampleRate);

        int ret = audioDecoder.start(codec, sampleRate, channel, sampleBytes);
        assert ret == 0;

        byte[] buffer = new byte[BUFF_LEN];
        if (codec == EAudioCodec.WAV) {
            stream.skip(44);
        }

        int read;
        while ((read = stream.read(buffer)) > 0) {
            splittedPcmData.write(audioDecoder.feed(buffer, read));
        }

        splittedPcmData.write(audioDecoder.stop());

        if (skipVad) {
            splittedPcmData.addOffset(0, splittedPcmData.length());
        } else {
            String json = VadSplitter.process(vadResPath, vadPauseTime, maxSpeechTime, splittedPcmData.getBuffer().array(), splittedPcmData.length());
            if (json != null) {
                JSONArray jsonArray = JSONArray.parseArray(json);
                for (int i = 0; i < jsonArray.size(); i++) {
                    JSONObject item = jsonArray.getJSONObject(i);
                    splittedPcmData.addOffset(item.getIntValue("start"), item.getIntValue("end"));
                }
            } else {
                splittedPcmData.addOffset(0, splittedPcmData.length());
            }
        }

        audioDecoder.destroy();

        return splittedPcmData;
    }
}
